package com.gmrz.uaf.protocol.v1.schema;

import com.gmrz.uaf.common.Constants;
import com.gmrz.uaf.protocol.v1.processor.exception.TLVParsingException;
import com.gmrz.uaf.protocol.v1.processor.exception.UAFErrorCode;
import com.gmrz.util.Convert;
import com.google.common.io.LittleEndianDataOutputStream;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.util.ArrayList;
import java.util.List;

public class KeyRegistrationData {
	private static final Logger LOG = LogManager
			.getLogger(KeyRegistrationData.class);

	private transient String aaid = null;
	private transient short authenticatorVersion = 0;
	private transient byte authenticatonMode = 0;
	private transient int keyFormat = 0;
	private transient int signAlgorithm = 0;
	private transient byte[] hashedFinalChallenge = null;
	private transient byte[] keyID = null;
	private transient int registerCounter = 0;
	private transient int signCounter = 0;
	private transient byte[] pubKey = null;
	private transient byte[] uvi = null;
	private transient List<byte[]> uvs = null;
	private transient byte[] preAuthCode = null;
	private transient byte[] refCode = null;
	private transient byte[] fidoPubkey = null;
	private transient String deviceID = null;

	private byte[] receivedKRDBytes = null;

	public void setReceivedKRDBytes(byte[] krdBytes) {
		this.receivedKRDBytes = krdBytes;
	}

	public byte[] getReceivedKRDBytes() {
		return this.receivedKRDBytes;
	}

	public String getAAID() {
		return this.aaid;
	}

	public void setAAID(String aaid) {
		this.aaid = aaid;
	}

	public short getAuthenticatorVersion() {
		return this.authenticatorVersion;
	}

	public void setAuthenticatorVersion(short authenticatorVersion) {
		this.authenticatorVersion = authenticatorVersion;
	}

	public byte getAuthenticatonMode() {
		return this.authenticatonMode;
	}

	public void setAuthenticatonMode(byte authenticatonMode) {
		this.authenticatonMode = authenticatonMode;
	}

	public int getKeyFormat() {
		return this.keyFormat;
	}

	public void setKeyFormat(int keyFormat) {
		this.keyFormat = keyFormat;
	}

	public int getSignAlgorithm() {
		return this.signAlgorithm;
	}

	public void setSignAlgorithm(int signAlgorithm) {
		this.signAlgorithm = signAlgorithm;
	}

	public byte[] getHashedFinalChallenge() {
		return this.hashedFinalChallenge;
	}

	public void setHashedFinalChallenge(byte[] challenge) {
		this.hashedFinalChallenge = challenge;
	}

	public byte[] getKeyID() {
		return this.keyID;
	}

	public void setKeyID(byte[] keyID) {
		this.keyID = keyID;
	}

	public int getRegisterCounter() {
		return this.registerCounter;
	}

	public void setRegisterCounter(int registerCounter) {
		this.registerCounter = registerCounter;
	}

	public int getSignCounter() {
		return this.signCounter;
	}

	public void setSignCounter(int signCounter) {
		this.signCounter = signCounter;
	}

	public byte[] getPubKey() {
		return this.pubKey;
	}

	public void setPubKey(byte[] pubKey) {
		this.pubKey = pubKey;
	}

	public byte[] getUvi() {
		return this.uvi;
	}

	public void setUvi(byte[] uvi) {
		this.uvi = uvi;
	}

	public List<byte[]> getUvs() {
		return uvs;
	}

	public void setUvs(List<byte[]> uvs) {
		this.uvs = uvs;
	}

	public byte[] getPreAuthCode() {
		return preAuthCode;
	}

	public void setPreAuthCode(byte[] preAuthCode) {
		this.preAuthCode = preAuthCode;
	}

	public byte[] getRefCode() {
		return refCode;
	}

	public void setRefCode(byte[] refCode) {
		this.refCode = refCode;
	}

	public byte[] getFidoPubkey() {
		return fidoPubkey;
	}

	public void setFidoPubkey(byte[] fidoPubkey) {
		this.fidoPubkey = fidoPubkey;
	}

	public String getDeviceID() {
		return deviceID;
	}

	public void setDeviceID(String deviceID) {
		this.deviceID = deviceID;
	}

	public String toString() {
		return "KeyRegistrationData [aaid=" + this.aaid
				+ ", authenticatorVersion=" + this.authenticatorVersion
				+ ", authenticatonMode=" + this.authenticatonMode
				+ ", keyFormat=" + this.keyFormat + ", signAlgorithm="
				+ this.signAlgorithm + ", hashedFinalChallenge="
				+ Convert.toBase64(this.hashedFinalChallenge) + ", keyID="
				+ Convert.toBase64(this.keyID) + ", registerCounter="
				+ this.registerCounter + ", signCounter=" + this.signCounter
				+ ", pubKey=" + Convert.toBase64(this.pubKey) + "]";
	}

	public byte[] toTLVBytes() throws TLVParsingException {
		ByteArrayOutputStream baos = new ByteArrayOutputStream();
		LittleEndianDataOutputStream ledo = new LittleEndianDataOutputStream(
				baos);
		try {
			byte[] aaidBytes = this.aaid == null ? new byte[0] : Convert
					.toUtf8(this.aaid);
			ledo.writeShort(Constants.UAFV1AssertionTag.TAG_AAID.getTagValue());
			ledo.writeShort(aaidBytes.length);
			ledo.write(aaidBytes);

			ledo.writeShort(Constants.UAFV1AssertionTag.TAG_ASSERTION_INFO
					.getTagValue());
			ledo.writeShort(7);
			ledo.writeShort(this.authenticatorVersion);
			ledo.writeByte(this.authenticatonMode);
			ledo.writeShort(this.signAlgorithm);
			ledo.writeShort(this.keyFormat);

			byte[] hashedFinalChallengeBytes = this.hashedFinalChallenge == null ? new byte[0]
					: this.hashedFinalChallenge;
			ledo.writeShort(Constants.UAFV1AssertionTag.TAG_FINAL_CHALLENGE
					.getTagValue());
			ledo.writeShort(hashedFinalChallengeBytes.length);
			ledo.write(hashedFinalChallengeBytes);

			byte[] keyIDBytes = this.keyID == null ? new byte[0] : this.keyID;
			ledo.writeShort(Constants.UAFV1AssertionTag.TAG_KEYID.getTagValue());
			ledo.writeShort(keyIDBytes.length);
			ledo.write(keyIDBytes);

			ledo.writeShort(Constants.UAFV1AssertionTag.TAG_COUNTERS
					.getTagValue());
			ledo.writeShort(8);
			ledo.writeInt(this.signCounter);
			ledo.writeInt(this.registerCounter);

			byte[] pubKeyBytes = this.pubKey == null ? new byte[0]
					: this.pubKey;
			ledo.writeShort(Constants.UAFV1AssertionTag.TAG_PUB_KEY
					.getTagValue());
			ledo.writeShort(pubKeyBytes.length);
			ledo.write(pubKeyBytes);

			if (this.uvi != null) {
				ledo.writeShort(Constants.UAFV1AssertionTag.AFI_USER_VERIFICATION_INDEX
						.getTagValue());
				ledo.writeShort(this.uvi.length);
				ledo.write(this.uvi);
			}

			if(this.uvs != null && this.uvs.size() > 0){
				ledo.writeShort(Constants.UAFV1AssertionTag.AFI_UVS_USER_VERIFICATION_SET.getTagValue());
				ledo.writeShort(this.uvs.size() * 32);
				for(byte[] uvi : this.uvs)
					ledo.write(uvi);
			}

			ledo.close();
		} catch (IOException ioe) {
			LOG.error("Failed to create TLV for KRD", ioe);
			throw new TLVParsingException(
					UAFErrorCode.PROTOCOL_TLV_PARSING_FAILED, ioe);
		}

		return baos.toByteArray();
	}

	public KeyRegistrationData parseTLVBytes(byte[] tlvBytes)
			throws TLVParsingException {
		ByteBuffer bb = ByteBuffer.wrap(tlvBytes)
				.order(ByteOrder.LITTLE_ENDIAN);
		try {
			while (bb.hasRemaining()) {
				short tag = bb.getShort();
				short len = bb.getShort();

				Constants.UAFV1AssertionTag t = Constants.UAFV1AssertionTag
						.getTagForValue(tag);
				byte[] val;
				switch (t.getTagValue()) {
					case 11787:
						val = new byte[len];
						bb.get(val);
						this.aaid = Convert.fromUtf8(val);
						break;
					case 11790:
						val = new byte[len];
						bb.get(val);
						ByteBuffer authInfoBuf = ByteBuffer.wrap(val).order(
							ByteOrder.LITTLE_ENDIAN);
						this.authenticatorVersion = authInfoBuf.getShort();
						this.authenticatonMode = authInfoBuf.get();
						this.signAlgorithm = authInfoBuf.getShort();
						if (val.length > 5)
							this.keyFormat = authInfoBuf.getShort();
						break;
					case 11786:
						val = new byte[len];
						bb.get(val);
						this.hashedFinalChallenge = val;
						break;
					case 11785:
						val = new byte[len];
						bb.get(val);
						this.keyID = val;
						break;
					case 11789:
						val = new byte[len];
						bb.get(val);
						ByteBuffer counterBuf = ByteBuffer.wrap(val).order(
							ByteOrder.LITTLE_ENDIAN);
						this.signCounter = counterBuf.getInt();
						if (counterBuf.hasRemaining())
							this.registerCounter = counterBuf.getInt();
						break;
					case 11788:
						val = new byte[len];
						bb.get(val);
						this.pubKey = val;
						break;
					case 259:
						val = new byte[len];
						bb.get(val);
						this.uvi = val;
						break;
					case 260:
						if (this.uvi == null) {
							val = new byte[len];
							bb.get(val);
							this.uvi = val;
						}
						break;

					case 261:
						val = new byte[len];
						LOG.info("uvs length=" + len);
						bb.get(val);
						if(len > 0 && (len % 32) == 0){
							this.uvs = new ArrayList<byte[]>(len/32);

							for(int i = 0; i < len/32; i++){
								byte[]  uvi = new byte[32];
								System.arraycopy(val, i*32, uvi, 0, 32);
								this.uvs.add(uvi);
							}
						}

						break;
					default:
					LOG.debug("Unrecognized tag in KRD[ " + tag + "]");
				}
			}
		} catch (Exception ex) {
			LOG.error("Error parsing Registration TLV data, ", ex);
			throw new TLVParsingException(
					UAFErrorCode.PROTOCOL_INVALID_REG_ASSERTION, ex);
		}
		return this;
	}
}
